package org.movee.jalgo.search;

import lombok.extern.slf4j.Slf4j;

import java.util.NoSuchElementException;
import java.util.Random;

/**
 * 实现跳表的基本功能
 *
 * @author 54740559@qq.com
 * @version 1.0
 *
 */
@Slf4j
public class SkipList<K extends Comparable<K>, V> {

    protected static final Random randomGenerator = new Random();
    protected static final double DEFAULT_PROBABILITY = 0.5;

    // 指向最高level的头节点
    private Node<K, V> head;
    // 新插入节点时，增加level的概率
    private final double probability;
    // 表中元素的数量
    private int size;

    public SkipList() {
        this(DEFAULT_PROBABILITY);
    }

    public SkipList(double probability) {
        this.head = new Node<K, V>(null, null, 0);
        this.probability = probability;
        this.size = 0;
    }

    public V get(K key) {
        checkKeyValid(key);
        Node<K, V> node = findNode(key);
        if (node.getKey().compareTo(key) == 0) {
            return node.getValue();
        } else {
            return null;
        }
    }

    public void add(K key, V value) {
        checkKeyValid(key);
        Node<K, V> node = findNode(key);
        if (node.getKey() != null && node.getKey().compareTo(key) == 0) {
            // 已经存在，将所有层的节点的value设置为新值
            node.setValue(value);
            while (node.down != null) {
                node = node.down;
                node.setValue(value);
            }
            return;
        }

        // 不存在，创建一个新节点。不存在时，level值永远是0
        Node<K, V> newNode = new Node<>(key, value, node.getLevel());
        horizontalInsert(node, newNode);

        int currentLevel = node.getLevel();
        int headLevel = head.getLevel();
        // 逐层判断是否需要在上一层增加节点
        while (needMoreLevel()) {
            if (currentLevel >= headLevel) {
                // 到了最高层，还需要增加level，head也增加一层
                Node<K, V> newHead = new Node<>(null, null, headLevel + 1);
                verticalLink(newHead, head);
                head = newHead;
                headLevel = head.getLevel();
            }

            // 寻找上一层的链表的前驱节点
            while (node.getUp() == null) {
                node = node.getPrev();
            }
            node = node.getUp();

            // 在上一层中插入新节点
            Node<K, V> tmp = new Node<>(key, value, node.getLevel());
            horizontalInsert(node, tmp);
            verticalLink(tmp, newNode);
            newNode = tmp;
            currentLevel++;
        }
        size++;
    }

    public void remove(K key) {
        checkKeyValid(key);
        Node<K, V> node = findNode(key);
        if (node == null || node.getKey().compareTo(key) != 0) {
            throw new NoSuchElementException("The key is not exist!");
        }

        // 逐层删除节点
        for (; node != null; node = node.getDown()) {
            Node<K, V> prev = node.getPrev();
            Node<K, V> next = node.getNext();
            // node的前驱肯定不为null
            prev.setNext(next);
            if (next != null) {
                next.setPrev(prev);
            }
        }

        // Adjust head，最高层没有其他节点了，层高降低一层
        while (head.getNext() == null && head.getDown() != null) {
            head = head.getDown();
            head.setUp(null);
        }
        size--;
    }

    public boolean contains(K key) {
        checkKeyValid(key);
        Node<K, V> node = findNode(key);
        return node.getKey().compareTo(key) == 0;
    }

    public int size() {
        return size;
    }

    public boolean empty() {
        return size == 0;
    }

    protected Node<K, V> findNode(K key) {
        Node<K, V> node = head;
        Node<K, V> next = null;
        Node<K, V> down = null;
        K nodeKey = null;

        while (true) {
            // 水平查找
            next = node.getNext();
            while (next != null && lessThanOrEqual(next.getKey(), key)) {
                node = next;
                next = node.getNext();
            }

            // 找到，直接返回
            nodeKey = node.getKey();
            if (nodeKey != null && nodeKey.compareTo(key) == 0) {
                break;
            }
            // 大于，继续到下一层查找
            down = node.getDown();
            if (down != null) {
                node = down;
            } else {
                break;
            }
        }

        return node;
    }

    protected void checkKeyValid(K key) {
        if (key == null) {
            throw new IllegalArgumentException("Key must be not null!");
        }
    }

    protected boolean lessThanOrEqual(K a, K b) {
        return a.compareTo(b) <= 0;
    }

    /**
     * 是否需要增加level
     * @return true：need；false：not need
     */
    protected boolean needMoreLevel() {
        return randomGenerator.nextDouble() < probability;
    }

    protected void horizontalInsert(Node<K, V> prev, Node<K, V> newNode) {
        newNode.setPrev(prev);
        newNode.setNext(prev.getNext());
        if (prev.getNext() != null) {
            prev.getNext().setPrev(newNode);
        }
        prev.setNext(newNode);
    }

    protected void verticalLink(Node<K, V> up, Node<K, V> down) {
        up.setDown(down);
        down.setUp(up);
    }

    protected static class Node<K extends Comparable<K>, V> {

        private K key;
        private V value;
        private int level;
        private Node<K, V> prev;
        private Node<K, V> next;
        private Node<K, V> up;
        private Node<K, V> down;

        public Node(K key, V value, int level) {
            this.key = key;
            this.value = value;
            this.level = level;
        }

        @Override
        public String toString() {
            StringBuilder sb = new StringBuilder();
            sb.append("Node[").append("key:");
            if (this.key == null) {
                sb.append("None");
            } else {
                sb.append(this.key);
            }

            sb.append(" value:");
            if (this.value == null) {
                sb.append("None");
            } else {
                sb.append(this.value);
            }
            sb.append("]");
            return sb.toString();
        }

        public K getKey() {
            return key;
        }

        public void setKey(K key) {
            this.key = key;
        }

        public V getValue() {
            return value;
        }

        public void setValue(V value) {
            this.value = value;
        }

        public int getLevel() {
            return level;
        }

        public void setLevel(int level) {
            this.level = level;
        }

        public Node<K, V> getUp() {
            return up;
        }

        public void setUp(Node<K, V> up) {
            this.up = up;
        }

        public Node<K, V> getDown() {
            return down;
        }

        public void setDown(Node<K, V> down) {
            this.down = down;
        }

        public Node<K, V> getNext() {
            return next;
        }

        public void setNext(Node<K, V> next) {
            this.next = next;
        }

        public Node<K, V> getPrev() {
            return prev;
        }

        public void setPrev(Node<K, V> prev) {
            this.prev = prev;
        }
    }

}
